3249. 统计好节点的数目

为保证权益，题目请参考 3249. 统计好节点的数目(From LeetCode).

解决方案1

说明

在这个问题中，首先是题目的理解，题目表意不清。参考题目理解如下：

翻译得依托，解释一下，这里是按照高度来区分父亲和孩子的，拿示例 1 举例：

• 根节点 0 的两个子节点分别为 1 和 2，两者包含的节点数都是 3，是好节点
• 节点 1 只有 两个 子节点 3 和 4，节点 0 是它爹不能算，所以是好节点
• 叶子节点没有子节点，也算节点数相同，是好节点 参考链接：点这里

这样的话，好节点也就是其所有的分支都具有同样的节点数量。

理解题意，可以认为题目的主要思想就是看如何从无向树中获取到以节点0为根的有向树的各个子节点是否满足好节点的定义。

这样，只需要通过DFS从0节点开始遍历每个节点即可。具体参考代码。

Python

from typing import List


class Solution:
    def countGoodNodes(self, edges: List[List[int]]) -> int:
        n = len(edges) + 1

        edge_map: dict[int, list[int]] = dict()
        for a, b in edges:
            if a not in edge_map:
                edge_map[a] = []
            edge_map[a].append(b)
            if b not in edge_map:
                edge_map[b] = []
            edge_map[b].append(a)

        visited = [False] * n
        self.ans = 0

        def dfs(node: int) -> int:
            visited[node] = True

            child_node_nums = []
            for child_node in edge_map[node]:
                if visited[child_node]:
                    continue
                child_node_nums.append(dfs(child_node))

            if len(child_node_nums) == 0 or len(set(child_node_nums)) == 1:
                self.ans += 1

            return sum(child_node_nums) + 1

        dfs(0)
        return self.ans